
local imageio = require 'imageio'
local alien = require 'alien'
local struct = require 'alien.struct'
local error, ipairs, setmetatable, assert, print, tostring, unpack
    = error, ipairs, setmetatable, assert, print, tostring, unpack
local cowrap, yield
    = coroutine.wrap, coroutine.yield
local strbyte, strchar
    = string.byte, string.char
local tconcat
    = table.concat

setfenv(1, imageio)

local CHAR = 'B'
local INT = 'i4'
local UINT = 'I4'
local SHORT = 'i2'
local USHORT = 'I2'

local filestruct_proto = {}
local filestruct_meta = {__index = filestruct_proto}

local function filestruct(def)
  local pattern = {'<'}
  for i, v in ipairs(def) do
    pattern[i+1] = v[2]
  end
  pattern = tconcat(pattern)
  def.pattern = pattern
  def.size = struct.size(pattern)
  return setmetatable(def, filestruct_meta)
end

function filestruct_proto:read(f)
  local size = self.size
  local data = f:read(size)
  if (data == nil) or (#data ~= size) then
    return nil
  end
  local values = {struct.unpack(self.pattern, data)}
  for i = 1,#self do
    values[self[i][1]] = values[i]
  end
  return values
end

function filestruct_proto:tostring(values)
  local arr = {}
  for i,v in ipairs(self) do
    arr[#arr+1] = assert(values[v[1]], "missing field: " .. v[1])
  end
  return struct.pack(self.pattern, unpack(arr))
end

local bmp_file_header = filestruct{
  {"magic_1", CHAR};
  {"magic_2", CHAR};
  {'file_size', UINT};
  {'reserved_1', USHORT};
  {'reserved_2', USHORT};
  {'bmp_offset', UINT};
}

local BITMAPINFOHEADER = filestruct{
  {'header_size', UINT};
  {'width', UINT};
  {'height', UINT};
  {'nplanes', USHORT};
  {'bpp', USHORT};
  {'compress_type', UINT};
  {'image_byte_size', UINT};
  {'hres', UINT};
  {'vres', UINT};
  {'ncolors', UINT};
  {'nimpcolors', UINT};
}

local BI_RGB = 0
local BI_RLE8 = 1 -- Can be used only with 8-bit/pixel bitmaps
local BI_RLE4 = 2 -- Can be used only with 4-bit/pixel bitmaps
local BI_BITFIELDS = 3 -- Can be used only with 16 and 32-bit/pixel bitmaps.
local BI_JPEG = 4 -- The bitmap contains a JPEG image
local BI_PNG = 5 -- The bitmap contains a PNG image

function read_value(f, valtype)
  valtype = '<' .. valtype
  local sz = struct.size(valtype)
  local data = f:read(sz)
  if (data == nil) or (#data ~= sz) then
    return nil
  end
  return (struct.unpack(valtype, data))
end

imageio.add_handler{
  pattern = 'BM';
  reader = function(f)
    local file_head = assert(bmp_file_header:read(f), 'unexpected end of file')
    local info_head_size = assert(read_value(f, UINT), 'unexpected end of file')
    f:seek("cur", -struct.size(UINT))
    local info_head
    if (info_head_size < 40) or (info_head_size >= 108) then
      error('unsupported info header size: ' .. tostring(info_head_size))
    else
      info_head = BITMAPINFOHEADER:read(f)
    end
    assert(info_head, 'unexpected end of file')
    local bpp = info_head.bpp
    local fmt
    local palette
    if (bpp == 16) then
      fmt = 'b5g5r5x1'
    elseif (bpp == 24) then
      fmt = 'b8g8r8'
    elseif (bpp == 32) then
      if (info_head.compress_type == BI_BITFIELDS) then
        local r, g, b = struct.unpack('>I4I4I4I4', f:read(4 * 4))
        if (r == 0x000000ff) and (g == 0x0000ff00) and (b == 0x00ff0000) then
          fmt = 'x8b8g8r8'
        else
          error(('unsupported bitfields: %08x %08x %08x'):format(r,g,b))
        end
      else
        fmt = 'b8g8r8x8'
      end
    else
      local bpp = bpp
      if (bpp ~= 1) and (bpp ~= 2) and (bpp ~= 4) and (bpp ~= 8) then
        error('unsupported bpp: ' .. tostring(bpp))
      end
      fmt = 'p' .. bpp
      palette = {}
      local ncolors = info_head.ncolors
      if (ncolors == 0) then
        ncolors = 2^bpp
      end
      assert(ncolors <= 2^bpp, 'palette too big for given bit depth')
      for i = 1, ncolors do
        local rgbx = f:read(4)
        palette[i] = {b = strbyte(rgbx,1)/255, g = strbyte(rgbx,2)/255, r = strbyte(rgbx,3)/255}
      end
    end
    local width = info_head.width
    local height = info_head.height
    local line_bytes = width * bpp/8
    local pitch = line_bytes + (4 - line_bytes) % 4
    -- jump to start of data
    f:seek('set', file_head.bmp_offset)
    -- return reader object
    return imageio.create_reader{
      bpp = bpp;
      format = fmt;
      pitch = pitch;
      
      width = width;
      height = height;
      palette = palette;
      
      getdata = cowrap(function()
        for i = 1, height do
          yield(height+1-i, 1, f:read(pitch))
        end
        return nil
      end);
    }
  end;
}
imageio.add_saver{
  name = 'bmp';
  extensions = {'bmp'};
  save = function(pixmap, f)
    local bpp, compress_type
    local palette
    local ncolors = 0
    local format = pixmap.format
    if (format == 'b8g8r8') then
      bpp = 24
      compress_type = BI_RGB
    elseif (format == 'p8') then
      bpp = 8
      compress_type = BI_RGB
      palette = assert(pixmap.palette, "no palette data found")
      ncolors = #palette
      local databuf = {}
      for i,v in ipairs(palette) do
        databuf[i] = strchar(v.b*255, v.g*255, v.r*255, 0)
      end
      palette = tconcat(databuf)
    elseif (format == 'p4') then
      bpp = 4
      compress_type = BI_RGB
      palette = assert(pixmap.palette, "no palette data found")
      ncolors = #palette
      local databuf = {}
      for i,v in ipairs(palette) do
        databuf[i] = strchar(v.b*255, v.g*255, v.r*255, 0)
      end
      palette = tconcat(databuf)
    elseif (format == 'p2') then
      bpp = 2
      compress_type = BI_RGB
      palette = assert(pixmap.palette, "no palette data found")
      ncolors = #palette
      local databuf = {}
      for i,v in ipairs(palette) do
        databuf[i] = strchar(v.b*255, v.g*255, v.r*255, 0)
      end
      palette = tconcat(databuf)
    elseif (format == 'p1') then
      bpp = 1
      compress_type = BI_RGB
      palette = assert(pixmap.palette, "no palette data found")
      ncolors = #palette
      local databuf = {}
      for i,v in ipairs(palette) do
        databuf[i] = strchar(v.b*255, v.g*255, v.r*255, 0)
      end
      palette = tconcat(databuf)
    elseif (format == 'x8b8g8r8') then
      bpp = 32
      compress_type = BI_BITFIELDS
      palette = "\000\000\000\255" .. "\000\000\255\000" .. "\000\255\000\000" .. "\000\000\000\000"
    else
      error('unsupported format: ' .. tostring(format))
    end
    palette = palette or ""
    
    local bmp_offset = bmp_file_header.size + BITMAPINFOHEADER.size + #palette
    local image_byte_size = pixmap.pitch * pixmap.height
    local file_size = bmp_offset + image_byte_size
    local palette_size = 3 * struct.size(UINT)
    f:write(bmp_file_header:tostring{
      magic_1 = strbyte('B');
      magic_2 = strbyte('M');
      file_size = file_size;
      reserved_1 = 0;
      reserved_2 = 0;
      bmp_offset = bmp_offset;
    })
    f:write(BITMAPINFOHEADER:tostring{
      header_size = BITMAPINFOHEADER.size;
      width = pixmap.width;
      height = pixmap.height;
      nplanes = 1;
      bpp = bpp;
      compress_type = compress_type;
      image_byte_size = image_byte_size;
      hres = 2835;
      vres = 2835;
      ncolors = ncolors;
      nimpcolors = 0;
    })
    f:write(palette)
    local row_pointers = pixmap:row_pointers()
    for i = pixmap.height, 1, -1 do
      f:write(alien.buffer(row_pointers[i]):tostring(pixmap.pitch))
    end
  end;
}
